# 导入需要的模块
import torch
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
from torch.utils.data import DataLoader

from model import Model
from dataset import HousePriceDataSet
from dataset import dataLoader

train_path = './dataset/train.csv'
test_path = './dataset/test.csv'

train_data = pd.read_csv(train_path)
test_data = pd.read_csv(test_path)

cols = ['OverallQual', 'GrLivArea', 'GarageCars', 'TotalBsmtSF', 'FullBath', 'TotRmsAbvGrd', 'YearBuilt']

# 得到训练集中salePrice的均值与标准差
train_salePrice = train_data['SalePrice']
salePrice_mu = np.mean(train_salePrice, axis=0)
salePrice_sigma = np.std(train_salePrice, axis=0)

# 加载训练数据集
dataset = HousePriceDataSet(root=train_path, features_name=cols, output_name='SalePrice')
train_loader = DataLoader(dataset=dataset, shuffle=False, batch_size=16)

# 加载测试集
test_feature = np.asarray([test_data[feature] for feature in cols])
test_feature_tensor = dataLoader(test_feature).T

model = Model()

criterion = torch.nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

epoch_list = []
loss_list = []


# 训练函数
def train(epoch):
    total_loss = 0
    for data in train_loader:
        x, y = data
        # 正向传播
        y_hat = model(x)
        # 计算损失
        loss = criterion(y_hat, y)

        total_loss += loss.item()

        # 反向传播
        loss.backward()
        # 权重更新
        optimizer.step()
        # 梯度清零
        optimizer.zero_grad()

    epoch_list.append(epoch)
    loss_list.append(total_loss / len(train_loader))

    # 如果训练到第300轮，学习率就变为之前的0.1倍
    if epoch % 300 == 0:
        for param_group in optimizer.param_groups:
            param_group['lr'] *= 0.1


if __name__ == '__main__':
    # 训练周期
    epochs = 1000

    for epoch in range(epochs):
        train(epoch)
    with torch.no_grad():
        y_pre = model(test_feature_tensor)

    y_pre = y_pre.squeeze(1)
    y_pre = y_pre.detach().numpy()

    my_pre = y_pre * salePrice_sigma + salePrice_mu

    Id = np.arange(1461, 1461 + len(my_pre))
    dataframe = pd.DataFrame({'Id': Id, 'SalePrice': my_pre})
    dataframe.to_csv(r"submission.csv", index=False, sep=',')

    plt.figure(1)
    plt.plot(loss_list)
    plt.ylabel('loss')
    plt.xlabel('epoch')

    other_pre_path = './dataset/sample_submission.csv'
    other_pre = pd.read_csv(other_pre_path)['SalePrice'].to_numpy()

    plt.figure(2)
    diff = abs(my_pre - other_pre).tolist()
    plt.plot(Id, diff)
    plt.ylabel('abs error/$')
    plt.xlabel('Id')

    plt.show()
